package com.noahsloan.nutils.streams;

import java.io.IOException;
import java.io.OutputStream;

/**
 * Writes to a series of output streams. When the count is reached with the
 * current stream, the stream is closed and the next is written to until its
 * limit is reached, etc.
 * 
 * @see #ChunkingOutputStream(OutputStream[], int[])
 * 
 * @author noah
 * 
 */
public class ChunkingOutputStream extends OutputStream {

    private final OutputStream[] streams;

    private final int[] counts;

    private int currentCount;

    private int current;

    /**
     * Writes the given number of bytes to the corresponding output streams. If
     * more bytes than sum(counts) are written to this stream, we wrap around
     * and start writing to the first one again.
     * <p>
     * NOTE that it will call close on a stream once it has reached the count,
     * so if the wrap around occurs, it will throw an IOException if the stream
     * cannot be written to again after closing.
     * <p>
     * A count less than 0 is considered to be infinity, so if you want extra
     * output to be discarded, make the last elements a NullOutputStream and -1
     * respectively.
     * 
     * @param streams
     * @param counts
     */
    public ChunkingOutputStream(OutputStream[] streams, int[] counts) {
        super();
        this.streams = new OutputStream[streams.length];
        System.arraycopy(streams, 0, this.streams, 0, streams.length);

        this.counts = new int[counts.length];
        System.arraycopy(counts, 0, this.counts, 0, counts.length);

        this.current = 0;
        this.currentCount = 0;
    }

    @Override
    public void write(int b) throws IOException {
        checkNext();
        streams[current].write(b);
        currentCount++;
    }

    /**
     * Determines if it is time to close the current stream and move on to the
     * next one.
     * 
     * @throws IOException
     */
    private void checkNext() throws IOException {
        if (counts[current] >= 0 && currentCount >= counts[current]) {
            streams[current].close();
            counts[current] = 0;
            current = (current + 1) % streams.length;
            currentCount = 0;
        }
    }

    @Override
    public void write(byte[] b) throws IOException {
        write(b, 0, b.length);
    }

    @Override
    public void write(final byte[] b, final int off, final int len)
            throws IOException {
        int offset = off;
        int end = len + off;
        while (offset < end) {
            checkNext();
            // figure out how many bytes we can write to current
            int count = Math.min(end - offset, counts[current] - currentCount);
            // write them
            streams[current].write(b, offset, count);
            // choke up
            offset += count;
        }
    }

}
